'''
@article{mahaur2022improved,
 title={Improved Residual Network Based on Norm-Preservation for Visual Recognition},
 author={Mahaur, Bharat and others},
 journal={Neural Networks},
 year={2022},
 publisher={Elsevier}
}

'''


import math
import torch
import torch.nn as nn
import torch.nn.functional as F
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# [u,v]表示二维DCT的分量下标，即实现对每一块采用不同的频率分量
# 这里u,v的值是根据论文中所做实验,选取正确率top_k得出
# 因为初始化的固定张量效果最好，所以这么写,开源代码里取top32,这里取top16
fidx_u_16 = [0,0,6,0,0,1,1,4,5,1,3,0,0,0,2,3] # 论文所选频率的水平指数
fidx_v_16 = [0,1,0,5,2,0,2,0,0,6,0,4,6,3,2,5] # 论文所选频率的垂直指数
# 获取一维通道的DCT
# L为输入的长度
def get_1d_dct(pos, freq, L):
    result = math.cos(math.pi * freq * (pos + 0.5) / L) / math.sqrt(L)
    if freq == 0:
        return result
    else:
        return result * math.sqrt(2)

# 推广到二维DCT
def get_dct_weights(input_width, input_height, input_channel, fidx_u= fidx_u_16, fidx_v= fidx_v_16):
    # 论文将频率空间分为7*7,实验结果表明,[u,v]=[0,0],即GAP的效果最好(乐),当然了，改进的点在于在信道注意力中使用多个频率分量
    # 使不同大小的频率与 7x7 频率空间相同
    fidx_horizon = [u*(input_width//7) for u in fidx_u]
    fidx_vertical = [v*(input_height//7) for v in fidx_v]
    dct_weights = torch.zeros(1, input_channel, input_width, input_height)
    c_part = input_channel // len(fidx_horizon)
    # 拆分通道
    for i, (u_x, v_y) in enumerate(zip(fidx_horizon, fidx_vertical)):
        for t_x in range(input_width):
            for t_y in range(input_height):
                dct_weights[:, i * c_part: (i+1)*c_part, t_x, t_y]\
                =get_1d_dct(t_x, u_x, input_width) * get_1d_dct(t_y, v_y, input_height)

    return dct_weights

# SE-Block单元--SEblock是一个子结构，几乎可以嵌入任何一个神经网络模型之中
class FcaBlock(nn.Module):
    def __init__(self, input_channel, width, height, reduction=16):
        super(FcaBlock, self).__init__()
        self.width = width
        self.height = height


        # self.adaptive_avg_pool = nn.AdaptiveAvgPool2d(1)  # 全局自适应池化
        self.register_buffer('pre_computed_dct_weights', get_dct_weights(self.width, self.height, input_channel))
        self.fc = nn.Sequential(
            nn.Linear(input_channel, input_channel // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(input_channel // reduction, input_channel),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, h, w = x.size()

        # squeeze操作:(b,c,h,w)->(b,c)
        #y = self.adaptive_avg_pool(x).view(b, c)
        y = F.adaptive_avg_pool2d(x, (self.height, self.width))
        y = torch.sum(y*self.pre_computed_dct_weights, dim=[2,3])

        # FC获取通道注意力权重，是具有全局信息的
        y = self.fc(y).view(b, c, 1, 1)

        # 注意力作用每一个通道上
        y = x * y.expand_as(x)
        # 残差连接
        return x+y


def conv3x3(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)


def conv1x1(in_planes, out_planes, stride=1):
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride,
                     bias=False)


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=None, downsampling=True,
                 downsampling_block=False, identity_blocks=False, bn0_exclude=False,width=32,height=32,reduction=16):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        # self.conv2 can be used to downsample if avgpool is not applied on stride != 1

        if not downsampling_block and not bn0_exclude:
            self.bn0 = norm_layer(inplanes)
        # self.downsample layers downsample the input when stride != 1

        self.conv1 = conv1x1(inplanes, planes)
        self.bn1 = norm_layer(planes)
        if not downsampling:
            self.conv2 = conv3x3(planes, planes, stride=1)  # only if using avgpool with stride != 1
        else:
            self.conv2 = conv3x3(planes, planes, stride)
        self.bn2 = norm_layer(planes)
        self.conv3 = conv1x1(planes, planes * self.expansion)
        self.avgpool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)

        if downsampling_block:
            self.bn3 = norm_layer(planes * self.expansion)

        if identity_blocks:
            self.bn3 = norm_layer(planes * self.expansion)

        self.fca = FcaBlock(planes, width, height, reduction)

        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

        self.downsampling = downsampling
        self.downsampling_block = downsampling_block
        self.identity_blocks = identity_blocks
        self.bn0_exclude = bn0_exclude

    def forward(self, x):
        identity = x

        if self.downsampling_block:
            out = self.conv1(x)  # downsampling block
        elif self.bn0_exclude:
            out = self.relu(x)  # start identity block
            out = self.conv1(out)
        else:
            out = self.bn0(x)  # middle identity block
            out = self.relu(out)
            out = self.conv1(out)

        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)

        if self.downsampling_block:
            out = self.bn3(out)



        if self.downsample is not None:
            identity = self.downsample(x)
            if not self.downsampling:
                out = self.avgpool(out)

        out = self.fca(out)
        out = out * out
        out += identity

        if self.identity_blocks:
            out = self.bn3(out)  # end identity block
            out = self.relu(out)

        return out


class LeNo_ResNet(nn.Module):

    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
                 norm_layer=None, dropout_prob0=0.0):
        super(LeNo_ResNet, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self.inplanes = 64
        self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, stride=1,
                                 padding=1, bias=False)
        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
                                 bias=False)  # input stem-B
        self.bn1 = norm_layer(64)
        self.relu = nn.ReLU(inplace=True)
        self.layer1 = self._make_layer(block, 64, layers[0], stride=2, norm_layer=norm_layer)
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, norm_layer=norm_layer)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, norm_layer=norm_layer)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))

        if dropout_prob0 > 0.0:
            self.dp = nn.Dropout(dropout_prob0, inplace=True)
            print("Using Dropout with the prob to set to 0 of: ", dropout_prob0)
        else:
            self.dp = None

        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)

    def _make_layer(self, block, planes, blocks, stride=1, norm_pooling=True, norm_layer=None):
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        downsample = None

        if stride != 1 and self.inplanes != planes * block.expansion:  # downsampling projection
            if not norm_pooling:
                downsample = nn.Sequential(
                    conv1x1(self.inplanes, planes * block.expansion),
                    norm_layer(planes * block.expansion),
                    nn.MaxPool2d(kernel_size=3, stride=stride, padding=1),
                )
            else:
                downsample = nn.Sequential(
                    nn.MaxPool2d(kernel_size=3, stride=stride, padding=1),
                    conv1x1(self.inplanes, planes * block.expansion),
                    norm_layer(planes * block.expansion),
                )
        elif self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion),
                norm_layer(planes * block.expansion),
            )
        elif stride != 1:
            downsample = nn.MaxPool2d(kernel_size=3, stride=stride, padding=1)

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample,
                            norm_layer, downsampling_block=True))
        self.inplanes = planes * block.expansion
        bn0_exclude = True
        for _ in range(1, (blocks - 1)):
            layers.append(block(self.inplanes, planes, norm_layer=norm_layer,
                                bn0_exclude=bn0_exclude))
            bn0_exclude = False

        layers.append(block(self.inplanes, planes, norm_layer=norm_layer,
                            identity_blocks=True, bn0_exclude=bn0_exclude))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1_1(x)
        x = self.conv1_2(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)

        if self.dp is not None:
            x = self.dp(x)

        x = self.fc(x)

        return x